#pragma region Copyright
/// <copyright>
/// Copyright (c) 2007 Ilia Shramko.
/// Released under the terms of the GNU Lesser General Public License.
/// </copyright>
#pragma endregion

#include "stdafx.h"

OSHProtocol::OSHProtocol() : _initialized(false)
{
	ATLTRACE_TS(atlTraceProtocol, TraceLevel::Verbose, "OSHProtocol::OSHProtocol\r\n");
}

OSHProtocol::~OSHProtocol()
{
	ATLTRACE_TS(atlTraceProtocol, TraceLevel::Verbose, "OSHProtocol::~OSHProtocol\r\n");
}


//*********************************************************
// IClassFactory
//*********************************************************
STDMETHODIMP OSHProtocol::CreateInstance(LPUNKNOWN pUnkOuter, REFIID riid, void** ppvObj)
{
	if(!_initialized)
		return E_UNEXPECTED;

	HRESULT hr = E_POINTER;
	if(ppvObj != NULL)
	{
		*ppvObj = NULL;

		if(pUnkOuter != NULL)
			return CLASS_E_NOAGGREGATION;

		CComPtr<IInternetProtocol> fileHelperProtocol;
		CComPtr<IInternetProtocol> httpHelperProtocol;

		hr = S_OK;
		if(_fileProtocolEnabled)
		{
			hr = CreateHelperInstance(IID_IInternetProtocol,
				_fileHelperClassFactory, reinterpret_cast<void**>(&fileHelperProtocol));
			if(FAILED(hr)) return hr;
		}

		if(_httpProtocolEnabled)
		{
			hr = CreateHelperInstance(IID_IInternetProtocol,
				_httpHelperClassFactory, reinterpret_cast<void**>(&httpHelperProtocol));
			if(FAILED(hr)) return hr;
		}

		CComObject<OSHProtocolHandler>* handler = NULL;
		hr = CComObject<OSHProtocolHandler>::CreateInstance(&handler);
		if(SUCCEEDED(hr))
		{
			hr = handler->QueryInterface(riid, ppvObj);
			if(SUCCEEDED(hr))
			{
				handler->Initialize(_basePath, _protocolSink, fileHelperProtocol, httpHelperProtocol);
			}
			else
			{
				delete handler;
			}
		}
	}
	return hr;
}

STDMETHODIMP OSHProtocol::LockServer(BOOL fLock)
{
	UNREFERENCED_PARAMETER(fLock);

	ATLTRACE_TS(atlTraceProtocol, TraceLevel::Verbose,
		"OSHProtocol::LockServer(%d)\r\n", fLock);

	if(!_initialized)
		return E_UNEXPECTED;

	return S_OK;
}


//*********************************************************
// IProtocol
//*********************************************************
STDMETHODIMP OSHProtocol::Register(VARIANT_BOOL enableHttpProtocol, VARIANT_BOOL enableFileProtocol,
	BSTR httpRootFolder, OSHLibrary::IProtocolSink* protocolSink)
{
	ATLTRACE_TS(atlTraceProtocol, TraceLevel::Verbose, "OSHProtocol::Register()\r\n");

	ATLASSERT(!_initialized);
	if(_initialized) return E_UNEXPECTED;

	//if(protocolSink == NULL) return E_INVALIDARG;

	_httpProtocolEnabled = (enableHttpProtocol == VARIANT_TRUE);
	_fileProtocolEnabled = (enableFileProtocol == VARIANT_TRUE);

	if(httpRootFolder != NULL)
		_basePath = httpRootFolder;

	_protocolSink = protocolSink;

	HRESULT hr = S_OK;
	CComPtr<IInternetSession> internetSession;
	hr = CoInternetGetSession(0, &internetSession, 0);
	ATLASSERT(SUCCEEDED(hr));
	if(FAILED(hr)) return hr;

	hr = CreateHelperClassFactory(CLSID_FileProtocol, _fileHelperClassFactory);
	if(SUCCEEDED(hr))
	{
		hr = CreateHelperClassFactory(CLSID_HttpProtocol, _httpHelperClassFactory);
		if(SUCCEEDED(hr))
		{
			IClassFactory* classFactory = static_cast<IClassFactory*>(this);
			hr = internetSession->RegisterNameSpace(classFactory,
				CLSID_OSHProtocolHandler, OSHNameSpace, 0, 0, 0);
			ATLASSERT(SUCCEEDED(hr));
			if(SUCCEEDED(hr))
			{
				ATLTRACE_TS(atlTraceProtocol, TraceLevel::Verbose,
					"OSH namespace handler has been registered.\r\n");

				_initialized = SUCCEEDED(hr);
				return S_OK;
			}
			// on error
			ReleaseHelperClassFactory(_httpHelperClassFactory);
		}
		ReleaseHelperClassFactory(_fileHelperClassFactory);
	}
	return hr;
}

STDMETHODIMP OSHProtocol::Unregister()
{
	ATLTRACE_TS(atlTraceProtocol, TraceLevel::Verbose, "OSHProtocol::Unregister()\r\n");

	//ATLASSERT(_initialized);
	if(!_initialized)
		return E_UNEXPECTED;

	HRESULT hr = S_OK;

	CComPtr<IInternetSession> internetSession;
	hr = CoInternetGetSession(0, &internetSession, 0);
	ATLASSERT(SUCCEEDED(hr));
	if(SUCCEEDED(hr))
	{
		IClassFactory* classFactory = static_cast<IClassFactory*>(this);
		hr = internetSession->UnregisterNameSpace(classFactory, OSHNameSpace);
		ATLASSERT(SUCCEEDED(hr));
		if(SUCCEEDED(hr))
		{
			hr = ReleaseHelperClassFactory(_httpHelperClassFactory);
			ATLASSERT(SUCCEEDED(hr));

			hr = ReleaseHelperClassFactory(_fileHelperClassFactory);
			ATLASSERT(SUCCEEDED(hr));

			ATLTRACE_TS(atlTraceProtocol, TraceLevel::Verbose,
				"OSH namespace handler has been unregistered.\r\n");

			_protocolSink = NULL;
			_basePath.Empty();
			_initialized = !SUCCEEDED(hr);
			return hr;
		}
	}

	return hr;
}


HRESULT OSHProtocol::CreateHelperClassFactory(REFCLSID helperClassID, CComPtr<IClassFactory>& classFactory)
{
	CComPtr<IClassFactory> helperClassFactory;
	HRESULT hr = CoGetClassObject(helperClassID, CLSCTX_ALL, 0, IID_IClassFactory,
		reinterpret_cast<void**>(&helperClassFactory));
	ATLASSERT(SUCCEEDED(hr) && helperClassFactory != 0);
	if(SUCCEEDED(hr) && helperClassFactory != 0)
	{
		hr = helperClassFactory->LockServer(TRUE);
		ATLASSERT(SUCCEEDED(hr));
		if(SUCCEEDED(hr))
		{
			ObjectLock lock(this);
			classFactory = helperClassFactory;
		}
	}
	return hr;
}

HRESULT OSHProtocol::ReleaseHelperClassFactory(CComPtr<IClassFactory>& classFactory)
{
	ATLASSERT(classFactory != NULL);

	HRESULT hr = classFactory->LockServer(FALSE);
	ATLASSERT(SUCCEEDED(hr));
	if(SUCCEEDED(hr))
	{
		ObjectLock lock(this);
		classFactory.Release();
	}
	return hr;
}

HRESULT OSHProtocol::CreateHelperInstance(REFIID riid, IClassFactory* classFactory, void** ppvObject)
{
	ATLASSERT(ppvObject != NULL);
	if(!ppvObject)
		return E_POINTER;

	*ppvObject = NULL;

	HRESULT hr = E_FAIL;

	ObjectLock lock(this);
	ATLASSERT(classFactory != NULL);
	if(classFactory != NULL)
	{
		hr = classFactory->CreateInstance(0, riid, ppvObject);
		ATLASSERT(SUCCEEDED(hr) && *ppvObject != NULL);	// unk can be NULL
		if(SUCCEEDED(hr) && *ppvObject == NULL)	// unk can be NULL if riid == IID_IUnknown
			hr = E_FAIL;
	}
	return hr;
}
